//
// Event.cxx
//
// Created by TB on 11/22/11.
// Copyright 2011 BNL. All rights reserved.
//

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <functional>
#include <iostream>

#include <TCollection.h> // For TIter
#include <TLorentzVector.h>
#include <TMath.h>
#include <TMCParticle.h>
#include <TObjArray.h>
#include <TPythia6.h>
#include <TRandom3.h>
#include <TRotation.h>
#include <TVector3.h>

#include <boost/foreach.hpp>

#include "tppmc/Distributions.h"
#include "tppmc/Enumerations.h"
#include "tppmc/Event.h"
#include "tppmc/FragmentationFunction.h"
#include "tppmc/Particle.h"
#include "tppmc/PartonContent.h"
#include "tppmc/RandomSpin.h"

namespace {
   bool isParton(int pdg) {
      return abs(pdg) < 7 or 21 == pdg;
   }

   // Helper functor for sorting Partons by group
   struct SortPartonsByGroup {
      bool operator()(const tppmc::Parton& a, const tppmc::Parton& b) {
         return a.Group() < b.Group();
      }
   };
}

namespace tppmc {

   Event::Event(double rotation, double etamin, double etamax, double minpt)
   : mRotation(rotation)
   , mEtaMin(etamin)
   , mEtaMax(etamax)
   , mMinPt(minpt)
     // Need full namespace specification in TClonesArray constructor.
   , mPartons("tppmc::Parton")
   , mHadrons("tppmc::Hadron") {
   }

   void Event::Update(Distributions& pdfs) {
      
      // Access the current PYTHIA state via the singleton TPythia6.
      TPythia6* pythia = TPythia6::Instance();
      
      mScatter = std::vector<Parton>(2, Parton());
      
      mProcessId = pythia->GetMSTI(1);
      mQSquared = pythia->GetPARI(22);
      mPartonPt = pythia->GetPARI(17);

      ComputeCentreOfMassEnergy();
      FillParticles(pdfs);
      
      Fragment(*(pdfs.Fragmentation()));
   } 

   void Event::ComputeCentreOfMassEnergy() {
      // Access the current PYTHIA state via the singleton TPythia6.
      TObjArray* particles = TPythia6::Instance()->ImportParticles();
      // Get the first 2 particles, which are the colliding beams
      TMCParticle* beam1 = dynamic_cast<TMCParticle*>(particles->At(0));
      TMCParticle* beam2 = dynamic_cast<TMCParticle*>(particles->At(1));
      // Use 4-vector class to compute centre-of-mass energy.
      TLorentzVector p1(beam1->GetPx(), beam1->GetPy(), beam1->GetPz(),
                        beam1->GetEnergy());
      TLorentzVector p2(beam2->GetPx(), beam2->GetPy(), beam2->GetPz(),
                        beam2->GetEnergy());
      mCentreOfMassEnergy = (p1 + p2).Mag();
   }

   void Event::FillParticles(Distributions& pdfs) {
      mPartons.Clear();
      mHadrons.Clear();
      // Access the current PYTHIA state via the singleton TPythia6.
      TObjArray* particles = TPythia6::Instance()->ImportParticles();
      // Loop over particles and extract partons and J/Psi.
      TIter next(particles);
      TMCParticle* particle(NULL);
      unsigned counter(0);
      // List to accumulate all partons.
      // We store them so we can sort them by group, so group 1 partons
      // precede group 2 partons, before adding them to the event.
      std::list<Parton> partons;
      while((particle = dynamic_cast<TMCParticle*>(next()))) {
         Int_t    type       = particle->GetKF();
         Int_t    status     = particle->GetKS();
         Int_t    parent     = particle->GetParent() - 1;
         Double_t px         = particle->GetPx();
         Double_t py         = particle->GetPy();
         Double_t pz         = particle->GetPz();
         // Event plane rotation.
         // mRotation was set externally before calling Update().
         Double_t pyth_pt  = sqrt(px * px + py * py); 
         Double_t pyth_phi = atan2(py, px);
         px = pyth_pt * cos(pyth_phi + mRotation);
         py = pyth_pt * sin(pyth_phi + mRotation);
         // Compute Sivers weights from the partons involved
         // in the initial hard interaction
         // First parton in hard interaction has index 4 in the particle list
         // Second parton in hard interaction has index 5 in the particle list
         if(4 == counter) {
            // this can go in a function
            mScatter.at(0) = Parton(px, py, pz, type, 1);
            mScatter.at(0).x = TPythia6::Instance()->GetPARI(33);
            mScatter.at(0).mSiversWeight = ComputeSiversWeight(mScatter.at(0), pdfs);
            // Note the minus sign
            int lhapdf = mScatter.at(0).Lhapdf();
            PartonContent values;
            pdfs.Unpolarised()->Evaluate(mScatter.at(0), mQSquared, values);
            mScatter.at(0).mF1 = values[lhapdf];
            pdfs.Sivers()->Evaluate(mScatter.at(0), mQSquared, values);
            mScatter.at(0).mSivers = values[lhapdf];
            pdfs.Transversity()->Evaluate(mScatter.at(0), mQSquared, values);
            mScatter.at(0).mH1 = values[lhapdf];
            // Generate proton spins for this event
            int spin = gRandom->Integer(2);
            if(spin == 0) {
               mScatter.at(0).pspin = -1; 
            } // if
            else {
               mScatter.at(0).pspin = 1;
            } // if
            mScatter.at(0).mSiversSpinFactor = ComputeSiversSpinFactor(mScatter.at(0));
            mScatter.at(0).spin =
            mScatter.at(0).pspin * ComputeTransverseSpin(mScatter.at(0), pdfs);
         } // if
         else if(5 == counter) {
            mScatter.at(1) = Parton(px, py, pz, type, 2);
            mScatter.at(1).x = TPythia6::Instance()->GetPARI(34);
            mScatter.at(1).mSiversWeight = ComputeSiversWeight(mScatter.at(1), pdfs);
            int lhapdf = mScatter.at(1).Lhapdf();
            PartonContent values;
            pdfs.Unpolarised()->Evaluate(mScatter.at(1), mQSquared, values);
            mScatter.at(1).mF1 = values[lhapdf];
            pdfs.Sivers()->Evaluate(mScatter.at(1), mQSquared, values);
            mScatter.at(1).mSivers = values[lhapdf];
            pdfs.Transversity()->Evaluate(mScatter.at(1), mQSquared, values);
            mScatter.at(1).mH1 = values[lhapdf];
            // Generate proton spins for this event
            int spin = gRandom->Integer(2);
            if(spin == 0) {
               mScatter.at(1).pspin = -1; 
            } // if
            else {
               mScatter.at(1).pspin = 1;
            } // if
            mScatter.at(1).mSiversSpinFactor = -ComputeSiversSpinFactor(mScatter.at(1));
            mScatter.at(1).spin =
            mScatter.at(1).pspin * ComputeTransverseSpin(mScatter.at(1), pdfs);
         } // else if
         
         if(isParton(type)) {
            if(6 == parent) { // From the first proton
               Parton parton(mScatter.at(0)); // Copy pdfs, proton spin, x
               // momentum and species can change after scattering, so update
               parton.SetPxPyPz(TVector3(px, py, pz));
               parton.type = type;
               partons.push_back(parton);
            } // if
            else if(7 == parent) { // From the second proton
               Parton parton(mScatter.at(1));
               // momentum and species can change after scattering, so update
               parton.SetPxPyPz(TVector3(px, py, pz));
               parton.type = type;
               partons.push_back(parton);
            } // else if
         } // if
         
         // Also accumulate J/Psi in the hadron list
         // \todo Add config flag option
         if((443 == type) and (1 == status)) {
            Hadron hadron;
            hadron.px =  px; 
            hadron.py =  py; 
            hadron.pz =  pz; 
            hadron.z = 1; 
            hadron.lasthad = 0; 
            hadron.type = JPSI; 
            hadron.mParton = Parton(0., 0., 0., 21, -1, 0, 0); // gluon 
            AddHadron(hadron);
         } // if
         
         ++counter;
      } // while
      partons.sort(SortPartonsByGroup());
      BOOST_FOREACH(const Parton& p, partons) {
         AddParton(p);
      } // BOOST_FOREACH
   }

   double Event::ComputeTransverseSpin(const Parton& parton, Distributions& pdfs) {
      RandomSpin random(pdfs.Unpolarised().get(),
                        pdfs.Transversity().get());
      return random.Generate(parton, mQSquared);
   }

   // Compute the ratio of the Sivers distribution to the unpolarised.
   double Event::ComputeSiversWeight(const Parton& parton, Distributions& pdfs) {
      PartonContent unpol;
      PartonContent sivers;
      pdfs.Unpolarised()->Evaluate(parton, mQSquared, unpol);
      pdfs.Sivers()->Evaluate(parton, mQSquared, sivers);
      int type = parton.Lhapdf();
      // Sivers is bounded by unpolarised
      double weight = sivers[type] / unpol[type];
      if(isnan(weight)) {
         weight = 0.;
      } // if
      return weight;
   }

   double Event::ComputeSiversSpinFactor(const Parton& parton) {
      if(parton.Pt() > 0.) {
         // SPECIAL NOTE: The sign of the Sivers spin factor has a minus sign relative to 
         // the extractions done by the Torino group because they used the Trento
         // convention to extract the Sivers distributions, where Sivers>0 corresponds
         // to the preference for a quark to move to the observers left if the proton is 
         // moving towards the observer and the spin is pointing up. pp measurements
         // use the opposite convention, and hence the minus sign. 
         /** \todo Reverse this, as I redefined the sivers such that it
         follows the pp convention i.e. sivers > 0 gives preference to
         observer's right*/
         return parton.pspin * (parton.px / parton.Pt());
      } // if
      else {
         return 0.;
      } // else
   }

   std::vector<tppmc::Parton*> Event::Partons() const {
      std::vector<Parton*> partons;
      TIter next(&mPartons);
      Parton* parton(NULL);
      while((parton = dynamic_cast<Parton*>(next()))) {
         partons.push_back(parton);
      } // while
      return partons;
   }

   std::vector<tppmc::Hadron*> Event::Hadrons() const {
      std::vector<Hadron*> hadrons;
      TIter next(&mHadrons);
      Hadron* hadron(NULL);
      while((hadron = dynamic_cast<Hadron*>(next()))) {
         hadrons.push_back(hadron);
      } // while
      return hadrons;
   }
   unsigned Event::AddParton(const Parton& parton) {
      const unsigned size = mPartons.GetEntries();
      new(mPartons[size]) Parton(parton);
      return mPartons.GetEntries();
   }

   unsigned Event::AddHadron(const Hadron& hadron) {
      const unsigned size = mHadrons.GetEntries();
      new(mHadrons[size]) Hadron(hadron);
      return mHadrons.GetEntries();
   }

   int Event::Fragment(FragmentationFunction& ff) {
      unsigned initialSize = mHadrons.GetEntries();
      BOOST_FOREACH(Parton* p, Partons()) {
         FragmentParton(*p, ff);
      } // BOOST_FOREACH
      return mHadrons.GetEntries() - initialSize;
   }

   int Event::FragmentParton(Parton& parton,
                             FragmentationFunction& ff) {
      unsigned initialSize = mHadrons.GetEntries();
      const double pTotalIntial = parton.PxPyPz().Mag();
      // We update the z range after each hadron is generated
      // so that we don't exceed the remaining momentum that
      // the parton has left.
      double zMax = ff.ZMax();
      // Also track the remaining momentum to calculate z.
      double pTotal = pTotalIntial;
      
      // Generate the rotation matrix that will put us into the 
      // fragmenting parton frame:
      TRotation rotM; 
      rotM.SetZAxis(parton.PxPyPz());
      TRotation rotMInv = rotM.Inverse();
      
      // We generate fragments until we're below a limit, so
      // keep a running sum of the total hadron momentum generated.
      Hadron runsum;
      runsum.px = 0.0;
      runsum.py = 0.0;
      runsum.pz = 0.0; 
      
      // Track the location of the spin axis: 
      const TVector3 spin_lab(0., double(parton.Spin()), 0.); // Y direction
      TVector3 spin_parton_frame = rotMInv * spin_lab; 
      
      // What is the angle between the spin axis in the parton frame
      // and the plane parallel to the parton momentum?
      // i.e. effectively no polarisation when the parton is parallel
      // to the spin direction.
      double cos_spin_depol = spin_parton_frame.Perp();
      parton.cos_spin_depol = cos_spin_depol;
      
      // Get the azimuthal angle of the spin axis in the new frame
      double phi_spin = TMath::ATan2(spin_parton_frame.Y(),
                                     spin_parton_frame.X());
      parton.phi_spin = phi_spin;
      if(parton.PxPyPz().Mag() > sqrt(ff.MinimumScale())) {
         do {
            // Use the parton pT^2 as the fragmentation scale.
            // If the parton pT is below the lower limit, use that.
            const double pt2 = std::max(pow(parton.Pt(), 2.),
                                        ff.MinimumScale());
            ff.SetZRange(ff.ZMin(), zMax);
            std::vector<Hadron> hadrons = ff.Fragment(parton, pt2);
            // Stop once we don't generate any more hadrons.
            if(hadrons.empty()) {
               break;
            } // if
            BOOST_FOREACH(Hadron& h, hadrons) {
               // Hadrons are in the parton frame.
               // Do some further calculations and rotate to lab frame.
               double newPz = h.Z() * pTotalIntial;
               pTotal -= newPz;
               double jt = h.PxPyPz().Pt();
               // Careful!
               // Phi we got from fragmentation is phi wrt spin direction.
               // The sense of the angle phi is opposite the sense
               // in the xy plane, hence the minus sign.
               // Also, the phi angle from fragmentation is the angle between
               // spin and (quark x hadron) momenta, so the hadron momentum is
               // at 90 degrees to this, hence additional pi/2 angle change.
               double phi = h.PxPyPz().Phi();
               h.collinsangle = phi;
               double new_px = jt * cos(phi_spin - TMath::PiOver2() - phi);
               double new_py = jt * sin(phi_spin - TMath::PiOver2() - phi);
               runsum.px += new_px;
               runsum.py += new_py;
               runsum.pz += newPz;
               
               TVector3 fragment(new_px, new_py, newPz);
               fragment.Transform(rotM); 
               
               h.px = fragment.X();
               h.py = fragment.Y();
               h.pz = fragment.Z();
               h.lasthad = 0; 
               if(h.PxPyPz().Eta() > mEtaMin and h.PxPyPz().Eta() < mEtaMax and
                  h.PxPyPz().Pt() > mMinPt) {
                  AddHadron(h);
               } // if
               zMax = pTotal / pTotalIntial; 
            } // BOOST_FOREACH
         } while(pTotal > 0.5 and zMax > ff.ZMin());
      } // if
      assert((pTotalIntial - pTotal - runsum.pz)<1.e-4);

      // Conserver any remaining parton momentum with a single soft hadron
      // Skip this if we aren't retaining the event (nhadrons == 0) or there
      // is an explicity eta cut (eta > -999).
      /**
       \todo We could potentially help charge non-conservation a little here by
       forcing the last hadron to have opposite charge to the sum of the other
       hadrons e.g if we have an excess of pi+ from the other fragmentations,
       force pi- and vice versa.
       \todo need to check momentum conservation works properly
      */
      if(pTotal > 0. and Hadrons().size() > 0 and mEtaMin<=-999. and mEtaMax>=999.) {
         // Running sum is in parton frame, transform difference 
         // back to lab frame - conserves momentum of fragmenting parton 
         TVector3 fragment(-1.0*runsum.px, -1.0*runsum.py, pTotal); 
         fragment.Transform(rotM); 
         Hadron hadron;
         hadron.px =  fragment.X();
         hadron.py =  fragment.Y();
         hadron.pz =  fragment.Z();
         hadron.z = pTotal / pTotalIntial;
         hadron.lasthad = 1;
         hadron.type = 111; // pi0
         hadron.mParton = parton;
         hadron.phigen = gRandom->Uniform(-TMath::Pi(), TMath::Pi());
         //calculate angle between hadron momentum and y-component of parton spin, projected into the plane perpendicular to the parton momentum
         TVector3 parton_mom = parton.PxPyPz();
         parton_mom = parton_mom.Unit();
         TVector3 spin_vec(0., parton.spin, 0.);
         TVector3 h_mom(hadron.px, hadron.py, hadron.pz);
         double projection_val = h_mom.Dot(parton_mom);
         TVector3 h_projection = h_mom - (projection_val * parton_mom);
         projection_val = spin_vec.Dot(parton_mom);
         TVector3 spin_projection = spin_vec - (projection_val * parton_mom);
         double cos_angle = (spin_projection.Unit()).Dot( h_projection.Unit() );
         hadron.collinsangle = acos(cos_angle);
         AddHadron(hadron);
      } // if
      return mHadrons.GetEntries() - initialSize;
   }

   Hadron Event::HadronSum(int group) const {
      Hadron sum;
      BOOST_FOREACH(Hadron* hadron, Hadrons()) {
         if(hadron->Group() == group or group == 0) {
            sum += *hadron;
         } // if
      } // BOOST_FOREACH
      return sum;
   }

   unsigned Event::Multiplicity(unsigned n) const {
      unsigned mult(0);
      BOOST_FOREACH(Hadron* hadron, Hadrons()) {
         if(unsigned(hadron->Group()) == n) ++mult;
      } // BOOST_FOREACH
      return mult;
   }

   //
   // class EventFactory
   //
   
   Event* EventFactory::Create(double rotation) {
      
      Event* event = new Event(rotation, -999., 999.);
      
      return event;
   }
} // namespace tppmc
